#!/usr/bin/env python3
import os
import sys
import argparse
import grpc
from time import sleep

sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections


def writeEcmp_groupRules(p4info_helper, sw, match_fields, action_params):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ecmp_group",
        match_fields={"hdr.ipv4.dstAddr": (match_fields[0], match_fields[1])},
        action_name="MyIngress.set_ecmp_select",
        action_params=action_params
    )
    sw.WriteTableEntry(table_entry)
    print(f"Installed ECMP group rule on {sw.name}")


def writeEcmp_nhopRules(p4info_helper, sw, index, action_params):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ecmp_nhop",
        match_fields={"meta.ecmp_select": index},
        action_name="MyIngress.set_nhop",
        action_params=action_params
    )
    sw.WriteTableEntry(table_entry)
    print(f"Installed ECMP nhop rule on {sw.name} for index {index}")


def writeIpv4_lpmRules(p4info_helper, sw, match_fields, action_params):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={"hdr.ipv4.dstAddr": (match_fields[0], match_fields[1])},
        action_name="MyIngress.ipv4_forward",
        action_params=action_params
    )
    sw.WriteTableEntry(table_entry)
    print(f"Installed IPv4 LPM rule on {sw.name} for {match_fields[0]}")


def main(p4info_file_path, bmv2_file_path):
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        switches = {}
        for i in range(1, 7):
            switches[f"s{i}"] = p4runtime_lib.bmv2.Bmv2SwitchConnection(
                name=f"s{i}",
                address=f"127.0.0.1:{50050 + i}",
                device_id=i - 1,
                proto_dump_file=f"logs/s{i}-p4runtime-requests.txt"
            )
            switches[f"s{i}"].MasterArbitrationUpdate()
            switches[f"s{i}"].SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                                          bmv2_json_file_path=bmv2_file_path)
            print(f"Installed P4 Program on s{i}")

        # ======= ECMP for s1 =======
        writeEcmp_groupRules(p4info_helper, switches["s1"], ["10.0.0.1", 32], {"ecmp_base": 0, "ecmp_count": 4})
        for i in range(4):
            writeEcmp_nhopRules(p4info_helper, switches["s1"], i, {
                "nhop_dmac": f"00:00:00:00:01:0{2+i}",
                "nhop_ipv4": "10.0.2.2",
                "port": 2 + i
            })
        writeIpv4_lpmRules(p4info_helper, switches["s1"], ["10.0.1.1", 32], {
            "dstAddr": "08:00:00:00:01:01",
            "port": 1
        })

        # ======= ECMP for s6 =======
        writeEcmp_groupRules(p4info_helper, switches["s6"], ["10.0.0.2", 32], {"ecmp_base": 0, "ecmp_count": 4})
        for i in range(4):
            writeEcmp_nhopRules(p4info_helper, switches["s6"], i, {
                "nhop_dmac": f"00:00:00:00:06:0{2+i}",
                "nhop_ipv4": "10.0.1.1",
                "port": 2 + i
            })
        writeIpv4_lpmRules(p4info_helper, switches["s6"], ["10.0.2.2", 32], {
            "dstAddr": "08:00:00:00:02:02",
            "port": 1
        })

        # ======= Rules for s2 ~ s5 =======
        for i in range(2, 6):
            writeIpv4_lpmRules(p4info_helper, switches[f"s{i}"], ["10.0.1.1", 32], {
                "dstAddr": "08:00:00:00:01:01",
                "port": 1
            })
            writeIpv4_lpmRules(p4info_helper, switches[f"s{i}"], ["10.0.2.2", 32], {
                "dstAddr": "08:00:00:00:02:02",
                "port": 2
            })

    except KeyboardInterrupt:
        print(" Shutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)
    finally:
        ShutdownAllSwitchConnections()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', type=str, default='./build/load_balance.p4.p4info.txt')
    parser.add_argument('--bmv2-json', type=str, default='./build/load_balance.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info) or not os.path.exists(args.bmv2_json):
        parser.print_help()
        print(f"\nMissing files: {args.p4info} or {args.bmv2_json}\n")
        parser.exit(1)

    main(args.p4info, args.bmv2_json)
